#!/bin/bash
# cuda=7
# trainer=new_fine_brats_epoch500_06
# task=Task607_CZ2
#
# CUDA_DEVICE_ORDER=PCI_BUS_ID \
# CUDA_VISIBLE_DEVICES=${cuda} \
# nnUNet_train \
# 3d_fullres \
# ${trainer} \
# ${task} \
# 0 \
# --npz \
# --fp32


nnUNet_dataset=/home/yusongli/Templates/yunet
export nnUNet_raw_data_base="${nnUNet_dataset}/nnUNet_raw"
export nnUNet_preprocessed="${nnUNet_dataset}/nnUNet_preprocessed"
export RESULTS_FOLDER="${nnUNet_dataset}/nnUNet_results"
export CUDA_DEVICE_ORDER=PCI_BUS_ID


TASK=203


function preprocess () {
    nnUNet_plan_and_preprocess -t "${TASK}"
}


function train () {
    CUDA="$1"
    FOLD="$2"
    # TRAINER=nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_7
    TRAINER="$3"

    CUDA_VISIBLE_DEVICES=${CUDA} \
    nnUNet_train \
        3d_fullres \
        "${TRAINER}" \
        "${TASK}" \
        "${FOLD}" \
        --npz
        # --fp32
}


function main () {
    # preprocess

    # FOLDS=(0 1 2 3 4)
    # for i in ${FOLDS[@]}; do
    #     train "${i}"
    # done

    # train 5 2 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet &
    # train 5 3 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet &
    # train 6 4 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet &
    #
    # train 6 0 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &
    # train 6 2 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &
    # train 7 3 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &
    # train 7 4 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &

    train 0 0 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &
    train 0 2 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_2 &
    train 0 4 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_3 &

    train 1 0 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_1 &
    train 1 2 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_2 &
    train 1 4 nnUNetTrainerV2BraTSRegions_DA4_BN_BD_res_axial_attention_unet_3 &
}


main
